#include <unistd.h>
#include <assert.h>
#include <poll.h>
#include <dlfcn.h>
#include <sys/epoll.h>
#include <sys/types.h>
#include <sys/socket.h>

#include <utility>

#include "../co_api.h"
#include "fd_reactor.h"
#include "hook_helper.h"
#include "hook_interface.h"
#include "../common/common_utils.h"
#include "../common/net_utils.h"

using namespace std;

int set_fd_block_raw(int fd)
{
    int flag = fcntl_f(fd, F_GETFL, 0);    
    if (flag == -1) {
        return -1;
    }
    flag &= ~O_NONBLOCK;
    return fcntl_f(fd, F_SETFL, flag);
}

int set_fd_unblock_raw(int fd)
{
    int flag = fcntl_f(fd, F_GETFL, 0);    
    if (flag == -1) {
        return -1;
    }
    return fcntl_f(fd, F_SETFL, flag | O_NONBLOCK);
}

BlockingGuard::BlockingGuard(const CoFdContextPtr& ptr)
{
    _ptr = ptr;
    int ret = set_fd_block_raw(_ptr->get_fd());
    assert(ret == 0);
}

BlockingGuard::~BlockingGuard()
{
    int ret = set_fd_unblock_raw(_ptr->get_fd());
    assert(ret == 0);
}


short poll_events_to_epoll_events(short events)
{
    // TODO: first finish read/write event, second finish other event 
    short new_events = 0;
    if (events & POLLIN) {
        new_events |= EPOLLIN; 
    }
    if (events & POLLOUT) {
        new_events |= EPOLLOUT; 
    }
    return new_events;
}

// #########################################################
//					    HookRealize	
// #########################################################
HookRealize::HookRealize()
{
    init();
}

void HookRealize::init()
{
	socket_f = (socket_t)dlsym(RTLD_NEXT, "socket");	
	close_f  = (close_t)dlsym(RTLD_NEXT,  "close");	

	connect_f = (connect_t)dlsym(RTLD_NEXT, "connect");	
	accept_f  = (accept_t)dlsym(RTLD_NEXT,  "accept");	

	write_f = (write_t)dlsym(RTLD_NEXT, "write");	
	read_f  = (read_t)dlsym(RTLD_NEXT,  "read");	

	send_f = (send_t)dlsym(RTLD_NEXT, "send");	
	recv_f = (recv_t)dlsym(RTLD_NEXT, "recv");	
    poll_f = (poll_t)dlsym(RTLD_NEXT, "poll");

    sleep_f  = (sleep_t)dlsym(RTLD_NEXT,  "sleep");
    usleep_f = (usleep_t)dlsym(RTLD_NEXT, "usleep");
    
    getsockopt_f = (getsockopt_t)dlsym(RTLD_NEXT, "getsockopt");
    setsockopt_f = (setsockopt_t)dlsym(RTLD_NEXT, "setsockopt");

    fcntl_f = (fcntl_t)dlsym(RTLD_NEXT, "fcntl");
    pipe_f  = (pipe_t)dlsym(RTLD_NEXT, "pipe");

    LOG_DEBUG("hook system api succeed");
}

int HookRealize::socket(int domain, int type, int protocol)
{
    int fd = socket_f(domain, type, protocol);
    LOG_DEBUG(
        "hook.socket, domain:%d, type:%d, protocol:%d, ret:%d",
        domain,
        type,
        protocol,
        fd
    );

    if (fd >= 0) {

        {
            int ret = set_fd_unblock_raw(fd);
            if (ret) {
                LOG_ERROR("set_fd_unblock_raw fail, fd:%d", fd);
            }
            assert(!ret);
        }

        {
            CoFdContext* ptr = new CoSocketContext(
                fd, 
                CO_FD_SOCKET, 
                {domain, type, protocol}
            );
            bool ret = g_manager.fd_context_mgr_ptr()->add(
                fd, 
                CoFdContextPtr(ptr)
            );
            assert(ret);
        }
    }
    return fd;
}
/*
int HookRealize::close(int fd)
{
	auto ctx = g_manager.fd_context_mgr_ptr()->get(fd);
	if (!ctx) {
		LOG_WARN("get ctx fail, fd:%d", fd);
		assert(ctx);
		errno = ENOENT;
		return -1;
	}

    if (!g_manager.fd_context_mgr_ptr()->del(fd)) {
		LOG_ERROR("del fail, fd:%d", fd);
		assert(false);
		errno = ENOENT;
		return -1;
    }

    if (!g_manager.fd_reactor_ptr()->detach(ctx)) {
		LOG_ERROR("fd_reactor detach fail, fd:%d", fd);	
		assert(false);
		errno = ENOENT;
        return -1;	
    }

    if (shutdown(fd, SHUT_RDWR) != 0)
    {
        LOG_WARN("shutdown fail, fd:%d", fd);
        return -1;
    }

    return 0;
}
*/
int HookRealize::close(int fd)
{
	auto ctx = g_manager.fd_context_mgr_ptr()->get(fd);
	if (!ctx) {
		LOG_WARN("get ctx fail, fd:%d", fd);
		assert(ctx);
		errno = ENOENT;
		return -1;
	}

    if (!g_manager.fd_context_mgr_ptr()->del(fd)) {
		LOG_ERROR("del fail, fd:%d", fd);
		assert(false);
		errno = ENOENT;
		return -1;
    }

    g_manager.fd_reactor_ptr()->detach(ctx, [this](FdContextPtr& ptr, bool is_release) {
        if (!is_release) {
           // if (shutdown(ptr->get_fd(), SHUT_RDWR) != 0) {
            if (shutdown(ptr->get_fd(), SHUT_WR) != 0) {
                LOG_WARN("shutdown fail, fd:%d, opt:SHUT_RDWR", ptr->get_fd());
            }
        }
    });
    return 0;
}

int HookRealize::connect(int fd, const struct sockaddr* addr, socklen_t addrlen)
{
    LOG_DEBUG("call function:connect, fd:%d", fd);

    auto ctx = g_manager.fd_context_mgr_ptr()->get(fd);
    if (!ctx) {
        LOG_ERROR("get ctx fail, fd:%d", fd);
    }
    assert(ctx);

    if (!is_in_co_env()) {
        LOG_DEBUG("system.connect, fd:%d", fd);
        BlockingGuard guard(ctx);
        return connect_f(fd, addr, addrlen);
    }

    int ret = connect_f(fd, addr, addrlen);
    LOG_DEBUG("connect, fd:%d, ret:%d, errno:%d", fd, ret, errno);

    if (!ret) {
        return 0; 
    } else if (errno != EINPROGRESS) {
        return ret; 
    } else if (!ctx->is_block()) {
        return ret;
    }

    int timeout = ctx->get_timeout(CO_FD_CONNECT_TIMEOUT);

    struct pollfd pfd[1];
    pfd[0].fd = fd;
    pfd[0].events = POLLOUT;
    ret = HookRealize::poll(pfd, 1, timeout);
    if (ret == 0) {
        errno = ETIMEDOUT;
        return -1;
    }

    int sock_err = 0;
    socklen_t len = sizeof(socklen_t);
    ret = getsockopt(fd, SOL_SOCKET, SO_ERROR, &sock_err, &len);
    if (ret) {
        LOG_ERROR("getsockopt fail, fd:%d", fd);
        return -1;
    }
    if (sock_err) {
        LOG_ERROR("getsockopt succeed, but has socket error, fd:%d, error:%d", fd, sock_err);
        errno = sock_err;
    }
    return !sock_err ? 0 : -1;
}

int HookRealize::accept(int fd, struct sockaddr* addr, socklen_t* addrlen)
{
    int cfd = read_write_common(
        fd, 
        "accept", 
        EPOLLIN, 
        CO_FD_READ_TIMEOUT, 
        accept_f, 
        addr, 
        addrlen
    );

    if (cfd >= 0) {

        int ret = set_fd_unblock_raw(fd);
        if (ret) {
            LOG_ERROR("set_fd_unblock_raw fail, fd:%d", fd);
        }
        assert(!ret);

        auto accept_ptr = g_manager.fd_context_mgr_ptr()->get(fd);
        if (!accept_ptr) {
            LOG_ERROR("get fd context ptr fail, fd:%d", fd);
        }
        assert(accept_ptr);

        auto client_ptr = new CoSocketContext(
            cfd, 
            CO_FD_SOCKET, 
            ((CoSocketContext*)accept_ptr.get())->get_attr()
        );
		LOG_DEBUG("add fd, ctx, fd:%d", cfd);
        g_manager.fd_context_mgr_ptr()->add(
            cfd, 
            CoFdContextPtr(client_ptr)
        );
    }
    return cfd;
}

int HookRealize::read(int fd, void* buf, size_t len)
{
    return read_write_common(
        fd, 
        "read", 
        EPOLLIN, 
        CO_FD_READ_TIMEOUT, 
        read_f, 
        buf, 
        len
    );
}

int HookRealize::write(int fd, const void* buf, size_t len)
{
    if (is_in_co_env()) {
        int ret = write_f(fd, buf, len);
        if (ret > 0 || (errno != EAGAIN && errno != EWOULDBLOCK)) {
            return ret;
        }
    }

    return read_write_common(
        fd,
        "write",
        EPOLLOUT,
        CO_FD_WRITE_TIMEOUT,
        write_f,
        buf,
        len
    );
}

int HookRealize::recv(int fd, void* buf, size_t len, int flags)
{
    return read_write_common(
        fd,
        "recv",
        EPOLLIN,
        CO_FD_READ_TIMEOUT,
        recv_f,
        buf,
        len,
        flags
    );
}

int HookRealize::send(int fd, const void* buf, size_t len, int flags)
{
    if (is_in_co_env()) {
        int ret = send_f(fd, buf, len, flags);
        if (ret > 0 || (errno != EAGAIN && errno != EWOULDBLOCK)) {
            return ret;
        }
    }

    return read_write_common(
        fd,
        "send",
        EPOLLOUT,
        CO_FD_WRITE_TIMEOUT,
        send_f,
        buf,
        len,
        flags
    );
}

int HookRealize::poll(struct pollfd* fds, nfds_t nfds, int timeout)
{
    auto co = g_manager.get_running_co();
    if (!co || !timeout) {
        LOG_DEBUG("use system.poll, nfds:%s, timeout:%d", nfds, timeout);
        return poll_f(fds, nfds, timeout); 
    }

    nfds_t invalids = 0;
    if (invalids == nfds) {
        LOG_ERROR("poll fds invalids, count:%d", nfds);
        errno == EINVAL;
        return -1;
    }

    int cid = co->_id;

    co->_state = SUSPEND;
    co->_suspend_state = SUSPEND_BY_IO_BLOCK;

	g_manager.switch_to_main([&, cid, nfds, timeout] {

        LOG_DEBUG("poll processing, nfds:%d", nfds);

		bool added = false;

        CoSuspendEntry entry;

        auto co = g_manager.get_co(cid);
        assert(co);

        co->_entry_flag = shared_ptr<CoEntryFlag>(new CoEntryFlag);

        entry.cid = co->_id;
        entry.ptr = co->_entry_flag;

	    for (nfds_t i = 0; i < nfds; i++) {

    	    pollfd& pfd = fds[i];
            LOG_DEBUG("poll processing, fd:%d", pfd.fd);
			if (pfd.fd < 0) {
                LOG_DEBUG("poll fd operate value invalid, fd:%d", pfd.fd);
				continue ;
			}
			pfd.revents = 0;

    	    auto ptr = g_manager.fd_context_mgr_ptr()->get(pfd.fd);
            if (!ptr) {
                LOG_ERROR("poll fd can`t find ptr, fd:%d", pfd.fd); 
            }
    	    assert(ptr);
           // LOG_DEBUG("poll fd:%d, ptr:%p, real_fd:%d", pfd.fd, ptr.get(), ptr->get_fd());

            short epoll_events = poll_events_to_epoll_events(pfd.events);
            if (!g_manager.fd_reactor_ptr()->set(ptr, epoll_events, &pfd.revents, entry)) {
				pfd.revents = POLLNVAL;
                LOG_DEBUG("poll fd operate POLLNVAL, fd:%d", pfd.fd);
				continue ;
			}
			added = true;
    	}
        LOG_DEBUG("poll added:%s", added ? "true" : "false");
		if (!added) {
			// no one set epoll succeed, awake suspend coroutine
            LOG_DEBUG("poll not add fd, resume coroutine, cid:%d", cid);
			g_manager.resume_co(cid);
			return ;
		} else if (timeout > 0) {
            LOG_DEBUG("poll add fd, set timer to resume coroutine, cid:%d, timer:%d", cid, timeout);
            g_manager.set_timer(
                timeout,
                [entry] {g_manager.resume_co_by_suspend_entry(entry);},
                CommonUtils::format_string("poll timer resume co, cid:%d, timer:%d", cid, timeout)
            );
		}
	});
    LOG_DEBUG("poll resumed, check fd event");

	int count = 0;
	for (nfds_t i = 0; i < nfds; i++) {
		if (fds[i].revents) {
			count++;
		}
	}
	return count;
}

unsigned int HookRealize::sleep(unsigned int seconds)
{
    if (!seconds) {
        return 0; 
    }
    if (!is_in_co_env()) {
        return sleep_f(seconds);
    }
    g_manager.sleep_ms(seconds * 1000);
    return 0;
}

int HookRealize::usleep(useconds_t usec, bool sleep_call)
{
	if (!usec) {
		return 0;
	}
	if (!is_in_co_env()) {
		return usleep_f(usec);
	}
	auto msec = usec / 1000;
	g_manager.sleep_ms(msec > 0 ? msec : 1);
	return 0;
}

int HookRealize::getsockopt(int fd, int level, int optname, void* optval, socklen_t* optlen)
{
	return getsockopt_f(fd, level, optname, optval, optlen);
}

int HookRealize::setsockopt(int fd, int level, int optname, const void* optval, socklen_t optlen)
{
	int ret = setsockopt_f(fd, level, optname, optval, optlen);
	if (ret == 0 && level == SOL_SOCKET) {
		if (optname == SO_RCVTIMEO || optname == SO_SNDTIMEO) {
			FdContextPtr ctx = g_manager.fd_context_mgr_ptr()->get(fd);
			if (ctx) {
				const timeval& tv = *(const timeval*)optval;
				int msec = tv.tv_sec * 1000 + tv.tv_usec / 1000;
				ctx->set_timeout(
					optname == SO_RCVTIMEO ? CO_FD_READ_TIMEOUT : CO_FD_WRITE_TIMEOUT,
					msec
				);
			}
		}
	}
	return ret;
}

int HookRealize::fcntl(int fd, int cmd)
{
    return fcntl_f(fd, cmd);
}

int HookRealize::fcntl(int fd, int cmd, int arg)
{
    int ret_value = 0;
    switch (cmd) {
        case F_DUPFD:
        case F_DUPFD_CLOEXEC:
            {
                int ret_value = fcntl_f(fd, cmd, arg);
                if (ret_value >= 0) {
                    auto fd_ptr = g_manager.fd_context_mgr_ptr()->get(fd);
                    assert(fd_ptr);

                    auto new_fd_ptr = fd_ptr->clone(ret_value);
                    auto ret = g_manager.fd_context_mgr_ptr()->add(
                        new_fd_ptr->get_fd(), 
                        new_fd_ptr
                    );
                    assert(ret);
                }
                break ;
            }
        case F_SETFL:
            {
                auto ptr = g_manager.fd_context_mgr_ptr()->get(fd);
                assert(ptr);
/*
                ret_value = fcntl_f(fd, cmd, arg); 
                if (!ret_value) {
                    ptr->set_unblock((arg & O_NONBLOCK) > 0);
                }
*/
                ptr->set_block_flag((arg & O_NONBLOCK) > 0);
                ret_value = fcntl_f(fd, cmd, arg | O_NONBLOCK); 
                break ;
            }
        default:
            {
                ret_value = fcntl_f(fd, cmd, arg);
                break ;
            }
    }
    return ret_value;
}

int HookRealize::fcntl(int fd, int cmd, struct flock* arg)
{
    return fcntl_f(fd, cmd, arg);
}

int HookRealize::fcntl(int fd, int cmd, struct f_owner_exlock* arg)
{
    return fcntl_f(fd, cmd, arg);
}

int HookRealize::pipe(int pipefd[2])
{
    int ret = pipe_f(pipefd);
    if (ret == 0) {

        CoFdContextPtr ptr1(new CoFdContext(pipefd[0], CO_FD_FILE));
        CoFdContextPtr ptr2(new CoFdContext(pipefd[1], CO_FD_FILE));

        auto ret = g_manager.fd_context_mgr_ptr()->add(pipefd[0], ptr1);
        assert(ret);

        ret = g_manager.fd_context_mgr_ptr()->add(pipefd[1], ptr2);
        assert(ret);
    }
    return ret;
}

// #########################################################
//					    FdContextMgr
// #########################################################
bool FdContextMgr::add(int fd, FdContextPtr ptr)
{
    LOG_DEBUG("add fd:%d", fd);
	lock_guard<mutex> lock(_mutex);
	return _fd_ctxs.insert(make_pair(fd, ptr)).second;
}

bool FdContextMgr::del(int fd)
{
    LOG_DEBUG("del fd:%d", fd);
	lock_guard<mutex> lock(_mutex);
	return _fd_ctxs.erase(fd) == 1 ? true : false;
}

FdContextPtr FdContextMgr::get(int fd)
{
	FdContextPtr ptr;
    LOG_DEBUG("get fd:%d", fd);
	lock_guard<mutex> lock(_mutex);
	auto iter = _fd_ctxs.find(fd);
	if (iter != _fd_ctxs.end()) {
		ptr = iter->second;
	}
	return ptr;
}

bool FdContextMgr::exists(int fd)
{
	lock_guard<mutex> lock(_mutex);
	return _fd_ctxs.find(fd) != _fd_ctxs.end() ? true : false;
}

